package com.yun.algorithmproblem.leetcode;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

public class Leetcode2368 {
    int cnt = 0;

    public int reachableNodes1(int n, int[][] edges, int[] restricted) {
        boolean[] isRestricted = new boolean[n];
        for (int x : restricted) {
            isRestricted[x] = true;
        }
        System.out.println(Arrays.toString(isRestricted));
        List<Integer>[] g = new List[n];
        for (int i = 0; i < n; i++) {
            g[i] = new ArrayList<>();
        }
        for (int[] v : edges) {
            g[v[0]].add(v[1]);
            g[v[1]].add(v[0]);
        }
        dfs(0, -1, isRestricted, g);
        return cnt;
    }

    private void dfs(int curr, int prev, boolean[] isRestricted, List<Integer>[] g) {
        cnt++;
        for (Integer next : g[curr]) {
            if (next != prev && !isRestricted[next]) {
                dfs(next, curr, isRestricted, g);
            }
        }
    }

    public int reachableNodes(int n, int[][] edges, int[] restricted) {
        boolean[] isRestricted = new boolean[n];
        for (int x : restricted) {
            isRestricted[x] = true;
        }

        UnionFind uf=new UnionFind(n);
        for (int[] v : edges) {
            if(isRestricted[v[0]]||isRestricted[v[1]]){
                continue;
            }
            uf.merge(v[0],v[1]);
        }
        return uf.count();
    }

    static class UnionFind {
        private int[] f;
        private int[] rank;

        public UnionFind(int n) {
            f = new int[n];
            rank = new int[n];
            for (int i = 0; i < n; i++) {
                f[i] = i;
            }
        }

        public void merge(int x, int y) {
            int rx = find(x);
            int ry = find(y);
            if (rx != ry) {
                if (rank[rx] > rank[ry]) {
                    f[ry] = rx;
                } else if (rank[rx] < rank[ry]) {
                    f[rx] = ry;
                } else {
                    f[ry] = rx;
                    rank[rx]++;
                }
            }
        }

        public int find(int x) {
            if (x != f[x]) {
                x = find(f[x]);
            }
            return f[x];
        }

        public int count() {
            int cnt = 0;
            int rt = find(0);
            for (int i = 0; i < f.length; i++) {
                if (rt == find(i)) {
                    cnt++;
                }
            }
            return cnt;
        }
    }

    public static void main(String[] args) {
        int n = 7;
        int[][] edges = {{0, 1}, {1, 2}, {3, 1}, {4, 0}, {0, 5}, {5, 6}};
        int[] restricted = {4, 5};
        Leetcode2368 test = new Leetcode2368();
        test.reachableNodes(n, edges, restricted);
    }
}
